import sys
import importlib
import marshal
import multiprocessing.shared_memory as shm
from importlib import abc

class SharedMemoryFinder(abc.MetaPathFinder):
    def __init__(self):
        self.shm_dict = {}
        self.shm_objects = {}  # 新增：用于保存创建的共享内存对象

    def find_spec(self, fullname, path, target=None):
        #print(f"Trying to find spec for: {fullname}")

        if fullname in sys.modules:
            return None
        
        if fullname in self.shm_dict:
            loader = SharedMemoryLoader(self.shm_dict[fullname])
            return importlib.util.spec_from_loader(fullname, loader)

        sys.meta_path.remove(self)
        spec = importlib.util.find_spec(fullname)
        sys.meta_path.insert(0, self)

        if spec and spec.loader and fullname not in self.shm_dict:
            code = spec.loader.get_code(fullname)
            if code:
                data = marshal.dumps(code)
                try:
                    shm_obj = shm.SharedMemory(create=True, size=len(data), name=fullname + "share114514")
                except FileExistsError:
                    shm_obj = shm.SharedMemory(name=fullname + "share114514")
                #print(f"Module {fullname} saved to shared memory with name: {shm_obj.name}")
                shm_obj.buf[:len(data)] = data
                # 尝试直接读取以确保它存在
                self.shm_dict[fullname] = (shm_obj.name, len(data))
                self.shm_objects[fullname] = shm_obj  # 保存共享内存对象
                loader = SharedMemoryLoader(self.shm_dict[fullname])
                spec.loader = loader

        
        return spec

class SharedMemoryLoader(abc.Loader):
    def __init__(self, shm_info):
        self.shm_info = shm_info

    def create_module(self, spec):
        return None

    def exec_module(self, module):
        shm_obj = shm.SharedMemory(name=self.shm_info[0])
        data = shm_obj.buf[:].tobytes()
        print("load from shared_memory:"+str(self.shm_info[0]))
        code = marshal.loads(data)
        exec(code, module.__dict__)

finder = SharedMemoryFinder()
sys.meta_path.insert(0, finder)
